#include "dag.h"
#include "dag_operators.h"
#include "dag_types.h"

using namespace dag;
using namespace std;

namespace {
	template <typename _type>
	void alloc_copy(size_t rows, size_t cols, 
					const _type* const* inVals,
					_type**& outVals) {
		outVals = new _type*[rows];
		for (size_t i = 0; i < rows; i++) {
			outVals[i] = new _type[cols];
			memcpy(outVals[i], inVals[i], sizeof(_type) * cols);
		}
	}
	
	template <typename _type>
	void dealloc(size_t rows, _type**& vals) throw() {
		try {
			if (vals == NULL) {
				return;
			}
			
			for (size_t i = 0; i < rows; i++) {
				delete[] vals[i];
				vals[i] = NULL;
			}
			delete vals;
			vals = NULL;
		}
		catch (...) {
			vals = NULL;
		}
	}
	
	template <typename _type>
	void serialize_matrix(ostream& bStream, 
						  const _type* const* mat,
						  size_t numRows,
						  size_t numCols) {
		dag_uint32_t num = numRows;
		bStream.write((const char*)&num, dag_uint32_size);
		num = numCols;
		bStream.write((const char*)&num, dag_uint32_size);
		
		for (size_t row = 0; row < numRows; row++) {
			bStream.write((const char*)mat[row], numCols * sizeof(_type));
		}
	}
	
	template <typename _type>
	void inflate_matrix(istream& bStream,
						_type**& mat,
						size_t& numRows,
						size_t& numCols) {
		mat = NULL;
		dag_uint32_t num = 0;
		bStream.read((char*)&num, dag_uint32_size);
		numRows = num;
		bStream.read((char*)&num, dag_uint32_size);
		numCols = num;
		
		try {
			if (numCols > 0) {
				mat = new _type*[numRows];
				memset(mat, NULL, numRows * sizeof(_type*));
				
				for (size_t row = 0; row < numRows; row++) {
					mat[row] = new _type[numCols];
					bStream.read((char*)mat[row], numCols * sizeof(_type));
				}
			}
		}
		catch (...) {
			if (mat != NULL) {
				for (size_t i = 0; i < numRows; i++) {
					delete[] mat[i];
					mat[i] = NULL;
				}
				delete[] mat;
				mat = NULL;
			}
			throw;
		}
	}
}

LinearOp::LinearOp()
:	m_InputSpec(FrameSpec()),
m_OutputSpec(FrameSpec()),
m_DTrans(NULL),
m_FTrans(NULL),
m_ITrans(NULL),
m_BTrans(NULL)
{
}

LinearOp::LinearOp(const LinearOp& copy) 
:	m_InputSpec(copy.m_InputSpec),
	m_OutputSpec(copy.m_OutputSpec),
	m_DTrans(NULL),
	m_FTrans(NULL),
	m_ITrans(NULL),
	m_BTrans(NULL)
{
	try {
		if (m_OutputSpec.bSize > 0) {
			alloc_copy(m_OutputSpec.bSize, 
					   m_InputSpec.bSize + 1,
					   copy.m_BTrans,
					   m_BTrans);
		}
		if (m_OutputSpec.iSize > 0) {
			alloc_copy(m_OutputSpec.iSize, 
					   m_InputSpec.iSize + 1,
					   copy.m_ITrans, 
					   m_ITrans);
		}
		if (m_OutputSpec.fSize > 0) {
			alloc_copy(m_OutputSpec.fSize, 
					   m_InputSpec.fSize + 1, 
					   copy.m_FTrans, 
					   m_FTrans);
		}
		if (m_OutputSpec.dSize > 0) {
			alloc_copy(m_OutputSpec.dSize, 
					   m_InputSpec.dSize + 1, 
	  				   copy.m_DTrans, 
					   m_DTrans);
		}
	}
	catch (...) {
		dealloc(m_OutputSpec.bSize, m_BTrans);
		dealloc(m_OutputSpec.iSize, m_ITrans);
		dealloc(m_OutputSpec.fSize, m_FTrans);
		dealloc(m_OutputSpec.dSize, m_DTrans);
		throw;
	}	
}

LinearOp::LinearOp(const FrameSpec& inSpec, const FrameSpec& outSpec,
				   const double* const* dTrans,
				   const float* const* fTrans,
				   const int* const* iTrans,
				   const bool* const* bTrans)
:	m_InputSpec(inSpec),
m_OutputSpec(outSpec),
m_DTrans(NULL),
m_FTrans(NULL),
m_ITrans(NULL),
m_BTrans(NULL) 
{
	try {
		if (outSpec.bSize > 0) {
			alloc_copy(outSpec.bSize, inSpec.bSize + 1, bTrans, m_BTrans);
		}
		if (outSpec.iSize > 0) {
			alloc_copy(outSpec.iSize, inSpec.iSize + 1, iTrans, m_ITrans);
		}
		if (outSpec.fSize > 0) {
			alloc_copy(outSpec.fSize, inSpec.fSize + 1, fTrans, m_FTrans);
		}
		if (outSpec.dSize > 0) {
			alloc_copy(outSpec.dSize, inSpec.dSize + 1, dTrans, m_DTrans);
		}
	}
	catch (...) {
		dealloc(outSpec.bSize, m_BTrans);
		dealloc(outSpec.iSize, m_ITrans);
		dealloc(outSpec.fSize, m_FTrans);
		dealloc(outSpec.dSize, m_DTrans);
		throw;
	}
}

LinearOp::~LinearOp() throw() {
	dealloc(m_OutputSpec.bSize, m_BTrans);
	dealloc(m_OutputSpec.iSize, m_ITrans);
	dealloc(m_OutputSpec.fSize, m_FTrans);
	dealloc(m_OutputSpec.dSize, m_DTrans);
}

const FrameSpec& LinearOp::inputSpec() const {
	return m_InputSpec;
}

const FrameSpec& LinearOp::outputSpec() const {
	return m_OutputSpec;
}

const double* const* LinearOp::doubleTransform() const {
	return m_DTrans;
}

const float* const* LinearOp::floatTransform() const {
	return m_FTrans;
}

const int* const* LinearOp::intTransform() const {
	return m_ITrans;
}

const bool* const* LinearOp::boolTransform() const {
	return m_BTrans;
}

auto_ptr<FrameOperator> LinearOp::inflate(istream& bStream) {
	auto_ptr<FrameSpec> inSpec(NULL);
	auto_ptr<FrameSpec> outSpec(NULL);
	inSpec = FrameSpec::inflate(bStream);
	outSpec = FrameSpec::inflate(bStream);
	
	auto_ptr<LinearOp> linOp(new LinearOp());
	
	inflate_matrix(bStream, linOp->m_BTrans, outSpec->bSize, inSpec->bSize);
	inflate_matrix(bStream, linOp->m_ITrans, outSpec->iSize, inSpec->iSize);
	inflate_matrix(bStream, linOp->m_FTrans, outSpec->fSize, inSpec->fSize);
	inflate_matrix(bStream, linOp->m_DTrans, outSpec->dSize, inSpec->dSize);
	linOp->m_InputSpec = *inSpec;
	linOp->m_OutputSpec = *outSpec;
	
	return auto_ptr<FrameOperator>(linOp.release());
}

void LinearOp::operate(const Frame& input, Frame& output) {
	if (static_cast<const FrameSpec&>(input) != 
		static_cast<const FrameSpec&>(m_InputSpec)) {
		throw UnequalFrameSpec(input, m_InputSpec);
	}
	
	if (static_cast<const FrameSpec&>(output) !=
		static_cast<const FrameSpec&>(m_OutputSpec)) {
		output.resize(m_OutputSpec);
	}
	
	if (m_InputSpec.bSize > 0 && m_OutputSpec.bSize > 0) {
		/* TODO: what to do about bools? */
		for (size_t i = 0; i < m_OutputSpec.bSize; i++) {
			output.bools[i] = true;
			for (size_t j = 0; j < m_InputSpec.bSize; j++) {
				output.bools[i] = output.bools[i] && 
				input.bools[j] && 
				m_BTrans[i][j];
			}
		}
	}
	
	if (m_InputSpec.iSize > 0 && m_OutputSpec.iSize > 0) {
		for (size_t i = 0; i < m_OutputSpec.iSize; i++) {
			output.ints[i] = 0;
			for (size_t j = 0; j < m_InputSpec.iSize; j++) {
				output.ints[i] += input.ints[j] * m_ITrans[i][j];
			}
		}
	}
	
	if (m_InputSpec.fSize > 0 && m_OutputSpec.fSize > 0) {
		for (size_t i = 0; i < m_OutputSpec.fSize; i++) {
			output.floats[i] = 0;
			for (size_t j = 0; j < m_InputSpec.fSize; j++) {
				output.floats[i] += input.floats[j] * m_FTrans[i][j];
			}
		}
	}
	
	if (m_InputSpec.dSize > 0 && m_OutputSpec.dSize > 0) {
		for (size_t i = 0; i < m_OutputSpec.dSize; i++) {
			output.doubles[i] = 0;
			for (size_t j = 0; j < m_InputSpec.dSize; j++) {
				output.doubles[i] += input.doubles[j] * m_DTrans[i][j];
			}
		}
	}
}

const string& LinearOp::_frameOperatorID() const {
	return OpId::linear_op;
}

void LinearOp::_serialize(ostream& binaryStream) const {
	m_InputSpec.serialize(binaryStream);
	m_OutputSpec.serialize(binaryStream);
	serialize_matrix(binaryStream, 
					 m_BTrans, 
					 m_OutputSpec.bSize,
					 m_InputSpec.bSize + 1);
	serialize_matrix(binaryStream, 
					 m_ITrans, 
					 m_OutputSpec.iSize,
					 m_InputSpec.iSize + 1);
	serialize_matrix(binaryStream, 
					 m_FTrans, 
					 m_OutputSpec.fSize,
					 m_InputSpec.fSize + 1);
	serialize_matrix(binaryStream, 
					 m_DTrans, 
					 m_OutputSpec.dSize,
					 m_InputSpec.dSize + 1);
}
